/*-------------------------------------------------------------------------*\
|                                                                           |
|   FILE:       THREADMANAGER.H                                             |
|   PURPOSE:    Contains a complete implementation of a simple thread		|
|				manager.													|
|   AUTHOR:     Bruce Belson                                                |
|   DATE:       24 May 2005													|
|   VERSION:    4.3.2 NT                                                    |
|   COPYRIGHT:  (c) Planatech Solutions Ltd, 2005							|
|                                                                           |
|   This software is the property of Planatech Solutions Ltd.  It is        |
|   licensed for the use of the purchaser only, and no license is granted   |
|   for the further distribution of this source code whether for sale or    |
|   otherwise, whether in its original form or amended. Any binary files    |
|   generated from this source code can be distributed for sale to third    |
|   parties only with the express written permission of Planatech Solutions |
|   Ltd. Planatech Solutions Ltd and their agents accept no liability for   |
|   losses incurred of any kind through the use of this software.           |
|                                                                           |
\*-------------------------------------------------------------------------*/

/*-------------------------------------------------------------------------*\
|																			|
|	USAGE:																	|
|																			|
|	All public methods of ThreadManager are safe to use from the main 		|
|	thread at any time. None of the other classes or methods are for		|
|	external use.															|
|																			|
|	A typical usage pattern is as follows:                                  |
|																			|
|	ThreadManager* mgr = new ThreadManager(...);                            |
|	if (!mgr->Init())                                                       |
|	{																		|
|		delete mgr;                                                         |
|		return;		// FAIL                                                 |
|	}																		|
|	while (...)	{                                                           |
|		...																	|
|		mgr->StartThread(...);                                              |
|		...																	|
|	}																		|
|	mgr->Term();                                                            |
|	delete mgr;                                                             |
|																			|
\*-------------------------------------------------------------------------*/

#ifndef __THREADMANAGER_H__
#define __THREADMANAGER_H__

#include <vector>
#include <set>

namespace threadmanager
{


/**
*	ThreadManager class
*
*	Manages a pool of threads, and assigns tasks to threads as available.
*/

class ThreadManager
{
public:
	
/**
*	Lockable class
*
*	Wraps a mutex. Used to share resources betweenn threads.
	*/

	class Lockable
	{
	protected:
		HANDLE m_mutex;
	public:
		Lockable()
		{
			m_mutex = ::CreateMutex(NULL, FALSE, NULL);
		}
		~Lockable()
		{
			if (m_mutex != NULL) ::CloseHandle(m_mutex);
		}
		bool Lock()
		{
			return ::WaitForSingleObject(m_mutex, INFINITE) == WAIT_OBJECT_0;
		}
		void Unlock()
		{
			::ReleaseMutex(m_mutex);
		}
};
	/**
	*	Defines a task passed to a worker thread. 
	*	Can be queued until a thread is available to run.
		*/

	typedef void (__stdcall *threadfn_type)(void*);
	class Task
	{
	public:
		ThreadManager* m_manager;	// Manager that created this task
		threadfn_type m_fn;			// User-defined function that will be called by this thread
		void* m_userData;			// User-defined data that will be passed to (*m_fn)()
		HANDLE m_handle;			// Thread handle (NULL if not currently running)
		int m_semaphoreIndex;		// If running, the index of the semaphore which the thread
		// will use to indicate completion

	public:
		Task(ThreadManager* manager, threadfn_type fn, void* userData)
				: m_manager(manager),
				m_fn(fn),
				m_userData(userData),
				m_handle(NULL),
				m_semaphoreIndex(-1)
		{}
	};

	/**
	*	Thread-safe collections of tasks
	*/

	class TaskQueue
	{
	protected:
		typedef std::vector<Task*> coll_type;
		coll_type m_coll;
		Lockable m_lock;
	public:
		void Push(Task* task, bool head = false)
		{
			if (m_lock.Lock())
			{
				if (head)
					m_coll.insert(m_coll.begin(), task);
				else
					m_coll.push_back(task);
				m_lock.Unlock();
			}
		}
		Task* Pop()
		{
			Task* task = 0;
			if (m_lock.Lock())
			{
				if (m_coll.size() > 0)
				{
					task = m_coll[0];
					m_coll.erase(m_coll.begin());
				}
				m_lock.Unlock();
			}
			return task;
		}
		void Clear()
		{
			if (m_lock.Lock())
			{
				for (coll_type::iterator it = m_coll.begin(); it != m_coll.end(); it++)
					delete (*it);
				m_coll.clear();
				m_lock.Unlock();
			}
		}
		size_t GetCount()
		{
			size_t count = 0;
			if (m_lock.Lock())
			{
				count = m_coll.size();
				m_lock.Unlock();
			}
			return count;
		}
	};
	class TaskSet
	{
	protected:
		typedef std::set<Task*> coll_type;
		coll_type m_coll;
		Lockable m_lock;
	public:
		void Add(Task* task)
		{
			if (m_lock.Lock())
			{
				m_coll.insert(task);
				m_lock.Unlock();
			}
		}
		void RemoveAll(std::vector<Task*>& tasks)
		{
			Task* task = 0;
			if (m_lock.Lock())
			{
				for (coll_type::iterator it = m_coll.begin(); it != m_coll.end(); it++)
					tasks.push_back(*it);
				m_coll.clear();
				m_lock.Unlock();
			}
		}
		void Remove(Task* task)
		{
			if (m_lock.Lock())
			{
				m_coll.erase(task);
				m_lock.Unlock();
			}
		}
		size_t GetCount()
		{
			size_t count = 0;
			if (m_lock.Lock())
			{
				count = m_coll.size();
				m_lock.Unlock();
			}
			return count;
		}
	};

public:
	/**
	*	Public methods
		*/

	/**
	*	Constructor
	*	Initializes a ThreadManager instance with the specified
	*	thread limit.
	*	Must be called from main thread.
	*/

	ThreadManager(size_t maxThreads)
			: m_maxThreads(maxThreads),
			m_scheduler(NULL),
			m_stoppedSemaphore(NULL)
	{
		m_semaphores = new HANDLE[m_maxThreads + 2];
		size_t i;
		for (i = 0; i < m_maxThreads + 2; i++)
			m_semaphores[i] = ::CreateSemaphore(NULL, 0, 1, NULL);
		m_semaphoreUsers = new Task*[m_maxThreads];
		for (i = 0; i < m_maxThreads; i++)
			m_semaphoreUsers[i] = 0;
		m_stoppedSemaphore = CreateSemaphore(NULL, 0, 1, NULL);
	}

	/**
	*	Destructor
	*	Must be called from main thread.
	*/

	virtual ~ThreadManager()
	{
		for (size_t i = 0; i < m_maxThreads; i++)
		{
			if (m_semaphores[i] != NULL)
				::CloseHandle(m_semaphores[i]);
		}
		::CloseHandle(m_stoppedSemaphore);
		delete[] m_semaphores;
		delete[] m_semaphoreUsers;
	}

	/**
	*	Prepares a ThreadManager instance for use.
	*	Returns false if initialization fails.
	*	Must be called from main thread.
	*/

	bool Init()
	{
		if (m_scheduler != NULL)
			return true;
		m_scheduler = ::CreateThread(NULL, 0, SchedulerThreadFn, this,
		                             0, NULL);
		return (m_scheduler != NULL);
	}

	/**
	*
	*	Terminates a ThreadManager.
	*	Stops all worker threads
	*	Must be called from main thread.
	*/

	void Term()
	{
		SignalScheduler(1);
		TRACE("Term: Waiting\n");
		::WaitForSingleObject(m_stoppedSemaphore, INFINITE);
		m_scheduler = NULL;
		TRACE("Term: Complete\n");
	}

	/**
	*	Creates a new task object, adds it to the queue,
	*	and signals to the scheduler that a new thread is required.
	*	Must be called from main thread.
	*/

	bool StartThread(threadfn_type fn, void* userData)
	{
		m_queue.Push(new Task(this, fn, userData));
		SignalScheduler(0);
		return true;
	}

	/**
	*	Returns the maximum number of worker threads available to
	*	the ThreadManager.
	*/

	size_t GetMaxThreadCount() const
	{
		return m_maxThreads;
	}

protected:
	/**
	*	Signals the scheduler thread by setting one of the semaphores
	*	for which the scheduler is waiting (thread-safe).
		*/

	void SignalScheduler(size_t index)
	{
		::ReleaseSemaphore(m_semaphores[index], 1, NULL);
	}

	/**
	*	Worker thread outer function
	*	Invokes user-defined function
	*	Adds itself to the list of dead threads
	*	Signals the scheduler that it has completed
	*/

	static unsigned long __stdcall WorkerThreadFn(void* pv)
	{
		Task* task = (Task*)pv;
		(*task->m_fn)(task->m_userData);
		task->m_manager->m_deadThreads.Add(task);
		task->m_manager->SignalScheduler(task->m_semaphoreIndex);
		return 0;
	}

	/**
	*	Scheduler thread outer function
	*/

	static unsigned long __stdcall SchedulerThreadFn(void* pv)
	{
		((ThreadManager*)pv)->SchedulerMain();
		return 0;
	}

	/**
	*	Scheduler thread
	*	Loops, waiting for a signal from any other thread
	*	Semaphore 0 indicates that the main thread has created a new task, 
	*	and added it to the queue
	*	Semaphore 1 indicates that the main thread has requested a close-down
	*	Other semaphores indicate that a worker thread has completed
	*/

	void SchedulerMain()
	{
		TRACE("SchedulerMain: Starting\n");
		while (true)
		{
			// Wait
			DWORD dwWait = ::WaitForMultipleObjects(m_maxThreads + 2, m_semaphores, FALSE, INFINITE);
			TRACE("SchedulerMain: WaitForMultipleObjects=%lu\n", dwWait);
			if (dwWait == WAIT_TIMEOUT)
			{
				continue;
			}
			else if ((dwWait >= WAIT_ABANDONED) && (dwWait < (WAIT_ABANDONED + m_maxThreads + 2)))
			{
				continue;
			}
			else if ((dwWait >= WAIT_OBJECT_0) && (dwWait < (WAIT_OBJECT_0 + m_maxThreads + 2)))
			{
				int index = (dwWait - WAIT_OBJECT_0);
				if (index == 0)
				{
					// New thread
					SchedulerStartThreads();
				}
				else if (index == 1)
				{
					// Close down
					SchedulerStopAllThreads();
					::ReleaseSemaphore(m_stoppedSemaphore, 1, NULL);
					TRACE("SchedulerMain: Exiting\n");
					return;
				}
				else
				{
					// Handle dead thread
					Task* dead = m_semaphoreUsers[index - 2];
					dead->m_semaphoreIndex = -1;
					m_semaphoreUsers[index - 2] = 0;
					m_deadThreads.Remove(dead);
					m_activeThreads.Remove(dead);
					delete dead;

					// Start more threads
					SchedulerStartThreads();
				}
			}

		}
	}

	/**
	*	If there any tasks in the queue, and there are any threads available,
	*	attempts to start as many tasks as possible.
	*/

	void SchedulerStartThreads()
	{
		while (m_activeThreads.GetCount() < m_maxThreads)
		{
			size_t index;
			for (index = 0; index < m_maxThreads; index++)
			{
				if (m_semaphoreUsers[index] == 0)
					break;
			}
			if (index == m_maxThreads)
				break;

			Task* task = m_queue.Pop();
			if (task == 0)
				break;

			task->m_handle = CreateThread(NULL, 0, WorkerThreadFn, task,
			                              CREATE_SUSPENDED, NULL);
			if (task->m_handle == NULL)
			{
				m_queue.Push(task, true);
				break;
			}
			task->m_semaphoreIndex = index + 2;
			m_semaphoreUsers[index] = task;
			m_activeThreads.Add(task);
			::ResumeThread(task->m_handle);
			TRACE("SchedulerStartThreads started thread %d: %lx\n", index, task->m_handle);
		}
	}

	/**
	*	Stops all worker threads, intrusively
	*/

	void SchedulerStopAllThreads()
	{
		for (size_t i = 0; i < m_maxThreads; i++)
		{
			Task* t = m_semaphoreUsers[i];

			if (t != 0 &&  t->m_handle != NULL)
			{
				::TerminateThread(t->m_handle, 0);
				::CloseHandle(t->m_handle);
				t->m_handle = NULL;
				m_semaphoreUsers[i] = 0;
			}
		}
		std::vector<Task*> tasks;
		m_activeThreads.RemoveAll(tasks);
		m_deadThreads.RemoveAll(tasks);
	}

protected:
	size_t			m_maxThreads;		// Maximum number of worker threads available to ThreadManager

	TaskQueue		m_queue;			// Queue of tasks waiting to be run (thread-safe)
	TaskSet			m_deadThreads;		// Worker treads that have completed (thread-safe)
	TaskSet			m_activeThreads;	// Worker threads that are running (thread-safe)

	HANDLE*			m_semaphores;		// Array of semaphores used to wake the scheduler
	// Item 0 is used to signal a new task waiting
	// Item 1 is used to signal close-down
	// Further items are used indicate completion of a worker thread
	Task**			m_semaphoreUsers;	// Array of worker threads, matching their semaphores
	HANDLE			m_stoppedSemaphore;	// Semaphore used to indicate that close-dwon is complete

	HANDLE			m_scheduler;		// Handle of scheduler thread
};

} // namespace

#endif //__THREADMANAGER_H__
